import math
import numpy as np
from typing import List
from ezgl import Matrix4x4


class BlenderInterface:
    """
    定义路点过渡方式
    """
    def distance(self, pose1, pose2):
        pass

    def blend_ratio(self, pose1, pose2):
        """
        混合比例 = 混合半径 / 两个位姿之间的距离
        """
        pass

    def interp(self, pose1, pose2, t):
        pass


class SegmentInterface:

    def update(self, poses, blender: BlenderInterface):
        pass

    def length(self):
        pass

    def get_pose(self, blender, t):
        pass


class PositionBlender(BlenderInterface):
    """
    根据位置距离进行插值
    """

    def __init__(self, bias_m):
        self._bias_m = bias_m

    def distance(self, pose1:Matrix4x4, pose2:Matrix4x4):
        return Matrix4x4.dist(pose1, pose2)

    def blend_ratio(self, pose1:Matrix4x4, pose2:Matrix4x4):
        dist = Matrix4x4.dist(pose1, pose2)
        if dist == 0:
            return 1
        return self._bias_m / dist

    def interp(self, pose1:Matrix4x4, pose2:Matrix4x4, t):
        return pose1.interp(pose2, t)


class PoseBlender(BlenderInterface):
    """
    同时考虑位置和姿态的插值
    """
    def __init__(self, bias_m, bias_rad=None):
        """ 1.5 rad ~ 1 m """
        self._bais_m = bias_m
        self._bais_rad = bias_rad if bias_rad is not None else bias_m * 1.5

    def distance(self, pose1:Matrix4x4, pose2:Matrix4x4):
        if not isinstance(pose2, Matrix4x4):
            pose2 = Matrix4x4.fromVector6d(*pose2)
        if not isinstance(pose1, Matrix4x4):
            pose1 = Matrix4x4.fromVector6d(*pose1)
        
        dist = Matrix4x4.dist(pose1, pose2)

        rad = np.arccos(np.clip(pose1.quat.dot(pose2.quat), -1, 1))
        if rad > math.pi / 2:
            rad = math.pi - rad
        return max(dist, rad/1.5)

    def blend_ratio(self, pose1:Matrix4x4, pose2:Matrix4x4):
        dist = self.distance(pose1, pose2)
        if dist == 0:
            return 1
        return min(self._bais_m / dist, self._bais_rad / (dist * 1.5))

    def interp(self, pose1:Matrix4x4, pose2:Matrix4x4, t):
        
        return pose1.interp(pose2, t)


class PoseSegment(SegmentInterface):

    def __init__(self, poses: List[Matrix4x4], blender: BlenderInterface):
        self.left_length = 0
        self.right_length = 0
        self.num_poses = 0  # 控制点的数量 2 or 3
        self.poses = []
        if len(poses) >= 2:
            self.update(poses, blender)

    def __repr__(self) -> str:
        return f"PoseSegment(left_len={self.left_length:.3f}, right_len={self.right_length:.3f}, num_poses={self.num_poses})"

    def update(self, poses, blender: BlenderInterface):
        self.poses = poses
        self.num_poses = len(poses)
        self.left_length = blender.distance(self.poses[0], self.poses[1])
        if self.num_poses == 3:
            self.right_length = blender.distance(self.poses[1], self.poses[2])

    def length(self):
        if self.num_poses == 2:
            return self.left_length
        return self.left_length + self.right_length

    def get_pose(self, blender, t):
        """
        计算二次贝塞尔曲线上的点, x in [0, 1]
        """
        if self.num_poses == 2:
            pose = blender.interp(self.poses[0], self.poses[1], t)
        elif self.num_poses == 3:
            if np.abs(self.left_length - self.right_length) < 1e-2:
                s = t
            else:  # 等距离时间标度, 使得 2 阶贝塞尔曲线速度均匀
                s = (np.sqrt(self.left_length**2 + (self.right_length**2-self.left_length**2) * t) -
                     self.left_length) / (self.right_length - self.left_length)

            pose1 = blender.interp(self.poses[0], self.poses[1], s)
            pose2 = blender.interp(self.poses[1], self.poses[2], s)
            pose = blender.interp(pose1, pose2, s)
        return pose


class PyCartPathMotion:
    """
    分段 6D 贝塞尔曲线
    """

    def __init__(self, poses: list, Blender: BlenderInterface, SegmentType):
        self.Segment = SegmentType
        self._total_len = 0
        self._segment_prefix_len = []
        self._blender = Blender
        self._segments: List[SegmentInterface] = []
        self._poses = []   # for visualization
        self.segment_poses(poses)

    def segment_poses(self, poses):
        """
        将控制点分成 line + bezier 的形式

        Parameters:
        - poses : list of poses
        """
        self._poses = []
        if len(poses) < 2:
            raise ValueError("At least 2 poses are required")

        last_pose = poses[0]
        for i in range(1, len(poses)):
            blend_ratio = self._blender.blend_ratio(poses[i-1], poses[i])
            blend_ratio = min(0.5, blend_ratio)
            pose1 = self._blender.interp(poses[i-1], poses[i], blend_ratio)
            pose2 = self._blender.interp(poses[i], poses[i-1], blend_ratio)
            self._segments.append(self.Segment([last_pose, poses[i-1], pose1], self._blender))
            self._segments.append(self.Segment([pose1, pose2], self._blender))
            last_pose = pose2
        self._segments.append(self.Segment([last_pose, poses[-1]], self._blender))

        for segment in self._segments:
            self._segment_prefix_len.append(self._total_len)
            self._total_len += segment.length()
            self._poses.extend(segment.poses)

    def get_pose(self, s: float):
        s = min(s, self._total_len)
        # 二分查找 s 在 self._segment_prefix_len 中的位置
        i = 0
        j = len(self._segment_prefix_len) - 1
        while i <= j:
            mid = i + (j-i) // 2
            if self._segment_prefix_len[mid] < s:
                i = mid + 1  # i 是第一个满足 self._segment_prefix_len[i] >= s 的位置
            else:
                j = mid - 1  # j 是最后一个满足 self._segment_prefix_len[j] < s 的位置
        i = max(0, i-1)
        if self._segments[i].length() == 0:
            t = 1
        else:
            t = (s - self._segment_prefix_len[i]) / self._segments[i].length()
        return self._segments[i].get_pose(self._blender, t)

    def curve(self, interval=0.1, length=None):
        if length is None:
            length = self._total_len

        s_val = np.arange(0, length+interval, interval)
        s_val[-1] = length

        curve = []
        for s in s_val:
            curve.append(self.get_pose(s))
        return curve
